package com.qzdatasoft.ratelimit.test;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Random;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.LongAdder;

/**
 * 独立压测工具（JDK 11+ HttpClient，可用 JDK21 运行）
 *
 * 运行示例：
 *   java -cp target/<your-app>.jar com.qzdatasoft.starter.ratelimit.LoadTester \
 *     http://localhost:8080/api/user/2 30 200
 *
 * 参数：
 *   args[0] targetUrl (默认 http://localhost:8080/)
 *   args[1] durationSeconds (压测时长，单位秒, 默认 30)
 *   args[2] concurrencyThreads (并发线程数, 默认 100)
 *   args[3] varyKeys true|false (是否随机 X-Forwarded-For / X-User，默认 true)
 *
 * 输出：
 *  压测总时长
 *  总请求数，RPS
 *  请求状态分布（200、429、其他）。
 *  延迟分布（P50、P90、P99、最大值）和 snapshot(延迟数据的只读副本，也是是响应时间的样本数据集合，用于安全、高效地统计性能指标)。
 */
public class LoadTester {
    public static void main(String[] args) throws Exception {
        String url = args.length > 0 ? args[0] : "http://localhost:8080/api/user/2";
        int durationSeconds = args.length > 1 ? Integer.parseInt(args[1]) : 30;
        int concurrency = args.length > 2 ? Integer.parseInt(args[2]) : 100;
        boolean varyKeys = args.length > 3 ? Boolean.parseBoolean(args[3]) : true;

        System.out.printf("LoadTester target=%s duration=%ds concurrency=%d varyKeys=%b%n",
                url, durationSeconds, concurrency, varyKeys);

        HttpClient client = HttpClient.newBuilder()
                .connectTimeout(Duration.ofSeconds(5))
                .version(HttpClient.Version.HTTP_1_1)
                .build();

        ExecutorService workers = Executors.newFixedThreadPool(concurrency);// 并发线程池
        CountDownLatch startLatch = new CountDownLatch(1);
        AtomicBoolean stopFlag = new AtomicBoolean(false);

        LongAdder totalReq = new LongAdder();
        LongAdder ok200 = new LongAdder();
        LongAdder tooMany = new LongAdder();
        LongAdder others = new LongAdder();

        List<Long> latencies = Collections.synchronizedList(new ArrayList<>());

        int perThreadVariants = Math.max(1, Math.min(1000, concurrency));// 每线程变体数，避免过多内存占用
        Random rnd = new Random();

        for (int i = 0; i < concurrency; i++) {
            final int tid = i;
            workers.submit(() -> {
                try {
                    startLatch.await();
                    while (!stopFlag.get()) {
                        String target = url;
                        HttpRequest.Builder rb = HttpRequest.newBuilder()
                                .uri(URI.create(target))
                                .timeout(Duration.ofSeconds(10))
                                .GET();

                        if (varyKeys) {
                            // 模拟不同 ip/user，通过 header 发送
                            String ip = "10.0." + (tid % 255) + "." + (rnd.nextInt(250) + 1);
                            String user = "user-" + (tid % perThreadVariants);
                            rb.header("X-Forwarded-For", ip);
                            rb.header("X-User", user);
                        }

                        HttpRequest req = rb.build();
                        long start = System.nanoTime();
                        try {
                            HttpResponse<Void> resp = client.send(req, HttpResponse.BodyHandlers.discarding());
                            long latencyMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start);
                            latencies.add(latencyMs);
                            totalReq.increment();

                            int status = resp.statusCode();
                            if (status == 200) ok200.increment();
                            else if (status == 429) tooMany.increment();
                            else {
                                System.out.printf("[%d] %s%n", status, req.uri());
                                others.increment();
                            }
                        } catch (Exception ex) {
                            System.out.printf("[Exception]%s: %s%n", req.uri(),ex.getMessage());
                            long latencyMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start);
                            latencies.add(latencyMs);
                            totalReq.increment();
                            others.increment();
                        }
                    }
                } catch (InterruptedException ignored) {
                }
            });
        }

        long startTime = System.nanoTime();
        startLatch.countDown();
        TimeUnit.SECONDS.sleep(durationSeconds);
        stopFlag.set(true);

        workers.shutdown();
        workers.awaitTermination(10, TimeUnit.SECONDS);

        long elapsedMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime);
        long total = totalReq.sum();
        long ok = ok200.sum();
        long t429 = tooMany.sum();
        long oth = others.sum();
        double rps = elapsedMs > 0 ? (total * 1000.0 / elapsedMs) : 0;

        List<Long> snapshot; // 统计快照
        synchronized (latencies) {
            snapshot = new ArrayList<>(latencies);
        }
        Collections.sort(snapshot);

        long p50 = percentile(snapshot, 50);
        long p90 = percentile(snapshot, 90);
        long p99 = percentile(snapshot, 99);
        long max = snapshot.isEmpty() ? 0 : snapshot.get(snapshot.size() - 1);
        long avgInterval = elapsedMs*1_000_000L / ok;
        long TPS = ok * 1000 / elapsedMs;
        System.out.println("======= Load Test Result =======");
        System.out.printf("Duration: %d ms%n", elapsedMs);
        System.out.printf("Total Requests: %d, RPS: %.2f%n", total, rps);
        System.out.printf("200: %d, 429: %d, others: %d%n", ok, t429, oth);
        System.out.printf("Latency ms p50=%d p90=%d p99=%d max=%d (n=%d)%n",
                p50, p90, p99, max, snapshot.size());
        System.out.printf("TPS：%d (based on successful requests)%n", TPS);
        System.out.printf("Current IntervalNanos：%d%n", avgInterval);
        System.out.printf("Expected IntervalNanos：query CustomRateLimitRule.intervalNanos from console%n");
        System.out.println("================================");

        // exit with success (0)
    }

    private static long percentile(List<Long> sorted, int perc) {
        if (sorted.isEmpty()) return 0;
        int idx = (int) Math.ceil(perc / 100.0 * sorted.size()) - 1;
        idx = Math.max(0, Math.min(sorted.size() - 1, idx));
        return sorted.get(idx);
    }
}

